Summary

Add FIPO (Future-KL Influenced Policy Optimization) as a built-in loss type, enabling dense token-level credit assignment for RL training without a value network.

Ma et al., "FIPO: Eliciting Deep Reasoning with Future-KL Influenced Policy Optimization" (arXiv:2603.19835)

Method

FIPO extends standard clipped PPO by re-weighting per-token advantages with Future-KL influence weights. While GRPO broadcasts a uniform sequence-level advantage to all tokens, FIPO modulates each token's contribution based on how the future trajectory evolves:

$$\text{FutureKL}_t = \sum_{k=t}^{T} M_k \cdot \gamma^{k-t} \cdot \Delta \log p_k$$

where $\Delta \log p_k = \log \pi_\theta(y_k \mid y_{<k}) - \log \pi_{\theta_{\text{old}}}(y_k \mid y_{<k})$ and $\gamma = 2^{-1/\tau}$ is an exponential decay factor.

The influence weight is computed as:

$$f_t = \text{clip}\left(\exp(\text{FutureKL}_t),; 1 - \epsilon_{f},; 1 + \epsilon_{f}\right)$$

The final FIPO loss modifies the clipped PPO objective with re-weighted advantages:

$$\mathcal{L}_t^{\text{FIPO}} = \min\left(r_t \cdot \hat{A}_t \cdot f_t,;; \text{clip}(r_t, 1{-}\epsilon, 1{+}\epsilon) \cdot \hat{A}_t \cdot f_t\right)$$

Key design choices

• influence_weights are detached from the computation graph — FIPO does not backpropagate through Future-KL
• Requires multi-step training per rollout (global_batch_size < rollout_batch_size × n_samples_per_prompt) so that $\theta$ drifts from $\theta_{\text{old}}$, making Future-KL non-trivial
• Uses chunked matrix multiplication for memory-efficient $O(B \cdot L^2)$ Future-KL computation

Changes

Usage

# Use with GRPO advantage estimator
--loss-type fipo_loss \
--advantage-estimator grpo \
--fipo-decay-rate 32.0 \
--fipo-chunk-size 128 \
--fipo-clip-ratio 0.2 \
--fipo-safety-thresh 3.0 \
--fipo-dual-clip-c 10.0 \
--global-batch-size 64   # must be < rollout_batch_size × n_samples_per_prompt

FIPO hyperparameter guide

Results

FIPO consistently increases raw reward (Qwen3.5-4B-Base)

A token-level importance-ratio heatmap under FIPO. The strong emphasis on “Alternative” suggests that FIPO reinforces branching points that trigger self-reflection and exploration of alternative solution paths, rather than treating all tokens uniformly.

cc. @zhuzilin